/*
 * HSM Proxy Project.
 * Copyright (C) 2013 FedICT.
 *
 * This is free software; you can redistribute it and/or modify it
 * under the terms of the GNU Lesser General Public License version
 * 3.0 as published by the Free Software Foundation.
 *
 * This software is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with this software; if not, see 
 * http://www.gnu.org/licenses/.
 */

package test.integ.be.fedict.hsm.ws;

import java.util.Set;
import java.util.UUID;

import javax.xml.namespace.QName;
import javax.xml.soap.SOAPMessage;
import javax.xml.soap.SOAPPart;
import javax.xml.ws.ProtocolException;
import javax.xml.ws.handler.MessageContext;
import javax.xml.ws.handler.soap.SOAPHandler;
import javax.xml.ws.handler.soap.SOAPMessageContext;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.ws.security.WSSecurityException;
import org.apache.ws.security.message.WSSecHeader;
import org.apache.ws.security.message.WSSecSAMLToken;
import org.apache.ws.security.message.WSSecTimestamp;
import org.apache.ws.security.saml.ext.AssertionWrapper;
import org.opensaml.DefaultBootstrap;
import org.opensaml.common.SAMLObject;
import org.opensaml.common.SAMLObjectBuilder;
import org.opensaml.saml1.core.Assertion;
import org.opensaml.xml.Configuration;
import org.opensaml.xml.ConfigurationException;
import org.opensaml.xml.XMLObjectBuilderFactory;
import org.opensaml.xml.io.Marshaller;
import org.opensaml.xml.io.MarshallingException;

public class WSS4JTestSOAPHandler implements SOAPHandler<SOAPMessageContext> {

	private static final Log LOG = LogFactory
			.getLog(WSS4JTestSOAPHandler.class);

	private static final XMLObjectBuilderFactory xmlObjectBuilderFactory;

	static {
		try {
			DefaultBootstrap.bootstrap();
		} catch (ConfigurationException e) {
			throw new RuntimeException(e);
		}
		xmlObjectBuilderFactory = Configuration.getBuilderFactory();
	}

	@Override
	public boolean handleMessage(SOAPMessageContext context) {
		Boolean outboundProperty = (Boolean) context
				.get(MessageContext.MESSAGE_OUTBOUND_PROPERTY);
		if (true == outboundProperty.booleanValue()) {
			try {
				handleOutboundMessage(context);
			} catch (Exception e) {
				LOG.debug("error adding WS-Security header: " + e.getMessage(),
						e);
				throw new ProtocolException(e);
			}
		}
		return true;
	}

	private void handleOutboundMessage(SOAPMessageContext context)
			throws WSSecurityException, ConfigurationException,
			MarshallingException {
		SOAPMessage soapMessage = context.getMessage();
		SOAPPart soapPart = soapMessage.getSOAPPart();

		WSSecHeader wsSecHeader = new WSSecHeader();
		wsSecHeader.insertSecurityHeader(soapPart);

		WSSecTimestamp wsSecTimeStamp = new WSSecTimestamp();
		wsSecTimeStamp.setTimeToLive(60);
		wsSecTimeStamp.build(soapPart, wsSecHeader);

		WSSecSAMLToken wsSecSAMLToken = new WSSecSAMLToken();

		Assertion assertion = buildObject(Assertion.DEFAULT_ELEMENT_NAME,
				Assertion.class);
		assertion.setID("assertion-" + UUID.randomUUID().toString());
		Marshaller marshaller = Configuration.getMarshallerFactory()
				.getMarshaller(assertion);
		marshaller.marshall(assertion);

		AssertionWrapper assertionWrapper = new AssertionWrapper(assertion);
		wsSecSAMLToken.build(soapPart, assertionWrapper, wsSecHeader);
	}

	private <T extends SAMLObject> T buildObject(QName name, Class<T> clazz) {
		SAMLObjectBuilder samlObjectBuilder = (SAMLObjectBuilder) xmlObjectBuilderFactory
				.getBuilder(name);
		return (T) samlObjectBuilder.buildObject();
	}

	@Override
	public boolean handleFault(SOAPMessageContext context) {
		return true;
	}

	@Override
	public void close(MessageContext context) {
	}

	@Override
	public Set<QName> getHeaders() {
		return null;
	}
}
